package com.xflib.shardingsphere.utils;


import org.apache.commons.lang3.StringUtils;
import org.apache.shardingsphere.infra.algorithm.keygen.core.KeyGenerateAlgorithm;
import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
import org.hibernate.HibernateException;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.hibernate.id.IdentifierGenerator;
import org.hibernate.metadata.ClassMetadata;

import java.io.Serializable;
import java.util.UUID;

public class SnowflakeGenerator implements IdentifierGenerator {
    private static final String PARAM_TYPE = "type";
    private static final String PARAM_TYPE_DEFAULT = "Long";

    private static KeyGenerateAlgorithm snowflakeShardingKeyGenerator;

    public static Comparable<?> getId(String type) {
        checkKeyGenerator();

        if (null != snowflakeShardingKeyGenerator) {
            try {
                Comparable<?> fileId = snowflakeShardingKeyGenerator.generateKeys(null, 1).stream().findFirst().get();
                return type == "String"? String.valueOf(fileId) : fileId;
            }catch (Exception e){/* do nothing */}
        }

        int abs = Math.abs(UUID.randomUUID().hashCode());
        return type == "String"? String.valueOf(abs) : Long.valueOf(abs);

    }

    @Override
    public Serializable generate(SharedSessionContractImplementor session, final Object obj)
            throws HibernateException {
        ClassMetadata metadata = session.getEntityPersister(null,obj).getClassMetadata();
        Serializable id = metadata.getIdentifier(obj, session);

        if(id != null && StringUtils.isNotBlank(id.toString())){
            return id;
        }

        String dataType = metadata.getIdentifierType().getReturnedClass().getSimpleName();

        return (Serializable) getId(dataType);
    }

    private static void checkKeyGenerator(){
        if (null == snowflakeShardingKeyGenerator) snowflakeShardingKeyGenerator =
        TypedSPILoader.findService(KeyGenerateAlgorithm.class, "SNOWFLAKE").orElse(null);
    }
}
